"""
将 MIDOG 2021 中  tif 转换成 yolo 系列需要的数据集

https://midog.deepmicroscopy.org/download-dataset/

m21
[1, 4, 7, 10, 13, 16, 19, 22, 25, 28, 31, 34, 37, 40, 43, 46, 49, 51, 54, 57, 60, 63, 66, 69, 72, 75, 78, 81, 84, 87, 90, 93, 96, 99, 101, 104, 107, 110, 113, 116, 119, 122, 125, 128, 131, 134, 137, 140, 143, 146, 149]
[2, 5, 8, 11, 14, 17, 20, 23, 26, 29, 32, 35, 38, 41, 44, 47, 50, 52, 55, 58, 61, 64, 67, 70, 73, 76, 79, 82, 85, 88, 91, 94, 97, 100, 102, 105, 108, 111, 114, 117, 120, 123, 126, 129, 132, 135, 138, 141, 144, 147, 150]
[3, 6, 9, 12, 15, 18, 21, 24, 27, 30, 33, 36, 39, 42, 45, 48, 53, 56, 59, 62, 65, 68, 71, 74, 77, 80, 83, 86, 89, 92, 95, 98, 103, 106, 109, 112, 115, 118, 121, 124, 127, 130, 133, 136, 139, 142, 145, 148]

f1:
    train:
        [2, 5, 8, 11, 14, 17, 20, 23, 26, 29, 32, 35, 38, 41, 44, 47, 50, 52, 55, 58, 61, 64, 67, 70, 73, 76, 79, 82, 85, 88, 91, 94, 97, 100, 102, 105, 108, 111, 114, 117, 120, 123, 126, 129, 132, 135, 138, 141, 144, 147, 150
        ,3, 6, 9, 12, 15, 18, 21, 24, 27, 30, 33, 36, 39, 42, 45, 48, 53, 56, 59, 62, 65, 68, 71, 74, 77, 80, 83, 86, 89, 92, 95, 98, 103, 106, 109, 112, 115, 118, 121, 124, 127, 130, 133, 136, 139, 142, 145, 148]

    test:
        [1, 4, 7, 10, 13, 16, 19, 22, 25, 28, 31, 34, 37, 40, 43, 46, 49, 51, 54, 57, 60, 63, 66, 69, 72, 75, 78, 81, 84, 87, 90, 93, 96, 99, 101, 104, 107, 110, 113, 116, 119, 122, 125, 128, 131, 134, 137, 140, 143, 146, 149]

f2:
    train:
        [1, 4, 7, 10, 13, 16, 19, 22, 25, 28, 31, 34, 37, 40, 43, 46, 49, 51, 54, 57, 60, 63, 66, 69, 72, 75, 78, 81, 84, 87, 90, 93, 96, 99, 101, 104, 107, 110, 113, 116, 119, 122, 125, 128, 131, 134, 137, 140, 143, 146, 149
        ,3, 6, 9, 12, 15, 18, 21, 24, 27, 30, 33, 36, 39, 42, 45, 48, 53, 56, 59, 62, 65, 68, 71, 74, 77, 80, 83, 86, 89, 92, 95, 98, 103, 106, 109, 112, 115, 118, 121, 124, 127, 130, 133, 136, 139, 142, 145, 148]

    test:
        [2, 5, 8, 11, 14, 17, 20, 23, 26, 29, 32, 35, 38, 41, 44, 47, 50, 52, 55, 58, 61, 64, 67, 70, 73, 76, 79, 82, 85, 88, 91, 94, 97, 100, 102, 105, 108, 111, 114, 117, 120, 123, 126, 129, 132, 135, 138, 141, 144, 147, 150]

f3:
    train:
        [1, 4, 7, 10, 13, 16, 19, 22, 25, 28, 31, 34, 37, 40, 43, 46, 49, 51, 54, 57, 60, 63, 66, 69, 72, 75, 78, 81, 84, 87, 90, 93, 96, 99, 101, 104, 107, 110, 113, 116, 119, 122, 125, 128, 131, 134, 137, 140, 143, 146, 149
        ,2, 5, 8, 11, 14, 17, 20, 23, 26, 29, 32, 35, 38, 41, 44, 47, 50, 52, 55, 58, 61, 64, 67, 70, 73, 76, 79, 82, 85, 88, 91, 94, 97, 100, 102, 105, 108, 111, 114, 117, 120, 123, 126, 129, 132, 135, 138, 141, 144, 147, 150]

    test:
        [3, 6, 9, 12, 15, 18, 21, 24, 27, 30, 33, 36, 39, 42, 45, 48, 53, 56, 59, 62, 65, 68, 71, 74, 77, 80, 83, 86, 89, 92, 95, 98, 103, 106, 109, 112, 115, 118, 121, 124, 127, 130, 133, 136, 139, 142, 145, 148]


"""
import json
import os
import numpy as np
from PIL import Image
import cv2
from tqdm import tqdm
from concurrent.futures import ThreadPoolExecutor

wsi_path = "/media/hsmy/wanghao_18T/dataset/MIDOG2021/tiff_sn/"
midog_json = "/media/hsmy/wanghao_18T/dataset/MIDOG2021/MIDOG.json"
data = json.load(open(midog_json))
png_path = "/media/hsmy/wanghao_18T/dataset/MIDOG2021/fold1/640_pn_/images/"
label_path = "/media/hsmy/wanghao_18T/dataset/MIDOG2021/fold1/640_pn_/labels/"
os.makedirs(png_path, exist_ok=True)
os.makedirs(label_path, exist_ok=True)

patch_size = 640  # 训练 image 大小
react_size = 50  # 框的大小 40X 下
react_size_percent = react_size / patch_size

filter_arr = [2, 5, 8, 11, 14, 17, 20, 23, 26, 29, 32, 35, 38, 41, 44, 47, 50, 52, 55, 58, 61, 64, 67, 70, 73, 76, 79, 82, 85, 88, 91, 94, 97, 100, 102, 105, 108, 111, 114, 117, 120, 123, 126, 129, 132, 135, 138, 141, 144, 147, 150
        ,3, 6, 9, 12, 15, 18, 21, 24, 27, 30, 33, 36, 39, 42, 45, 48, 53, 56, 59, 62, 65, 68, 71, 74, 77, 80, 83, 86, 89, 92, 95, 98, 103, 106, 109, 112, 115, 118, 121, 124, 127, 130, 133, 136, 139, 142, 145, 148]

save_gt = False
if save_gt:
    gt_path = f"/media/hsmy/wanghao_18T/dataset/MIDOG2021/fold1/640_pn_/gt/"
    os.makedirs(gt_path, exist_ok=True)

def do_convert(wsi_file):
    image_id = int(wsi_file[:3])
    if image_id not in filter_arr:
        return
    file_name = wsi_file.split('.')[0]
    file_path = os.path.join(wsi_path, wsi_file)
    img = cv2.imread(file_path)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    height, width, _ = img.shape
    mask = np.zeros((height, width)).astype(np.uint8)

    annotations = [anno for anno in data["annotations"] if
                   anno["image_id"] == image_id]
    if len(annotations) > 0:
        for anno in annotations:
            bbox_arr = anno["bbox"]
            x = bbox_arr[1]
            y = bbox_arr[0]
            x = int(x + 25)
            y = int(y + 25)
            mask[x, y] = anno["category_id"]

    overlap = 0
    for h in range(0, height, patch_size - overlap):
        h_start = h
        h_end = h + patch_size
        if h_end > height:
            h_start = height - patch_size
            h_end = height
        for w in range(0, width, patch_size - overlap):
            w_start = w
            w_end = w + patch_size
            if w_end > width:
                w_start = width - patch_size
                w_end = width

            # 切割块
            png_patch = img[h_start:h_end, w_start:w_end]
            mask_patch = mask[h_start:h_end, w_start:w_end]
            gt_patch = None

            png_name = f"{file_name}_{w_start}_{h_start}.png"
            txt_name = f"{file_name}_{w_start}_{h_start}.txt"

            if mask_patch.max() == 0:
                Image.fromarray(png_patch).save(png_path + png_name)
            else:
                if save_gt:
                    gt_patch = np.asarray(png_patch)
                # mitosis
                points = np.argwhere(mask_patch == 1)
                for point in points:
                    label = (
                        0,
                        point[1] / patch_size,
                        point[0] / patch_size,
                        react_size_percent,
                        react_size_percent
                    )
                    with open(label_path + txt_name, 'a') as f:
                        f.write(('%g ' * len(label)).rstrip() % label + '\n')

                    if save_gt:
                        offset = 25
                        gt_patch = cv2.rectangle(
                            gt_patch,
                            (point[1] - offset, point[0] - offset),
                            (point[1] + offset, point[0] + offset),
                            (255, 0, 0), 2, cv2.LINE_AA)

                # hard n
                points = np.argwhere(mask_patch == 2)
                for point in points:
                    label = (
                        1,
                        point[1] / patch_size,
                        point[0] / patch_size,
                        react_size_percent,
                        react_size_percent
                    )
                    with open(label_path + txt_name, 'a') as f:
                        f.write(('%g ' * len(label)).rstrip() % label + '\n')

                    if save_gt:
                        offset = 25
                        # 黄色 hard n
                        gt_patch = cv2.rectangle(
                            gt_patch,
                            (point[1] - offset, point[0] - offset),
                            (point[1] + offset, point[0] + offset),
                            (255, 255, 0), 2, cv2.LINE_AA)

                Image.fromarray(png_patch).save(png_path + png_name)
                if save_gt:
                    Image.fromarray(gt_patch).save(gt_path + png_name)
    print(f"{wsi_file} done")


with ThreadPoolExecutor(max_workers=20) as executor:
    for wsi_file in os.listdir(wsi_path):
        executor.submit(do_convert, wsi_file)


# for wsi_file in sorted(os.listdir(wsi_path)):
#     do_convert(wsi_file)